"""
    @file Sklearn
    @description Decision Tree By Sklearn
    @author Synhard
    @tel 13001321080
    @id 21126338
    @email 823436512@qq.com
    @date 2021-09-25 19:21
    @version 1.0
"""
# 构建模型
import seaborn as sns
from matplotlib import pyplot as plt
from sklearn.tree import DecisionTreeClassifier, export_graphviz
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.model_selection import train_test_split, GridSearchCV
import pandas as pd

decisionTree = DecisionTreeClassifier()
source = pd.read_excel("heart.xlsx")
label = source['target']
data = source.drop('target', axis=1)
train_X, test_X, train_y, test_y = train_test_split(data, label, random_state=3, test_size=0.3)

a = source.columns[0:-1]
# 训练数据
decisionTree = DecisionTreeClassifier(max_depth=5, criterion="entropy")
# decisionTree = DecisionTreeClassifier(max_depth=5, criterion="gini")
decisionTree.fit(train_X, train_y)

# 预测数据
pred_y = decisionTree.predict(test_X)

# 评估模型
print(accuracy_score(test_y, pred_y))

print(decisionTree.score(train_X, train_y))

print(decisionTree.score(test_X, test_y))

param_test = {'max_features': ['auto', 'sqrt', 'log2'],
              'min_samples_split': list(range(2, 20)),
              'min_samples_leaf': list(range(1, 12))

              }

tree_gv = GridSearchCV(estimator=decisionTree, param_grid=param_test, cv=5)
tree_gv.fit(train_X, train_y)

# 最优参数

# 预测数据
pred_y = tree_gv.predict(test_X)
print(classification_report(test_y, pred_y))

# 可视化决策树
export_graphviz(
    decisionTree,
    out_file="D:\研一\AI\dt\dt_sklearn_entropy.dot",
    feature_names=source.columns[0:-1],
    rounded=True,
    filled=True
)

# export_graphviz(
#     decisionTree,
#     out_file="D:\研一\AI\dt\dt_sklearn_gini.dot",
#     feature_names=source.columns[0:-1],
#     rounded=True,
#     filled=True
# )

fig1 = plt.figure(figsize=(3 * 5, 1 * 4))
matrix = pd.DataFrame(confusion_matrix(test_y, pred_y))
sns.heatmap(matrix, annot=True, cmap='OrRd')
plt.title('Confusion Matrix -- %s ')
plt.show()
